import os
import warnings
import pandas as pd
import numpy as np
from sklearn.impute import KNNImputer
from sklearn.preprocessing import StandardScaler, LabelEncoder, OneHotEncoder
from sklearn.feature_selection import SelectKBest, f_classif, mutual_info_classif, RFE
from sklearn.ensemble import RandomForestClassifier
from sklearn.decomposition import PCA
from sklearn.feature_selection import VarianceThreshold
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats

# 设置忽略警告
warnings.filterwarnings("ignore")

# 创建输出目录
output_dir = "output"
analysis_dir = "analysis_results"
os.makedirs(output_dir, exist_ok=True)
os.makedirs(analysis_dir, exist_ok=True)

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False


# 定义KNN填充函数（带标准化）
def knn_impute(df, numeric_cols, n_neighbors=5):
    """使用KNN填充数值型缺失值，包含标准化处理"""
    if df[numeric_cols].isnull().sum().sum() == 0:
        return df

    # 复制数据避免修改原始数据
    df_filled = df.copy()

    # 标准化数据（KNN对尺度敏感）
    scaler = StandardScaler()
    scaled_data = scaler.fit_transform(df_filled[numeric_cols])

    # 应用KNN填充
    imputer = KNNImputer(n_neighbors=n_neighbors)
    filled_data = imputer.fit_transform(scaled_data)

    # 逆标准化
    df_filled[numeric_cols] = scaler.inverse_transform(filled_data)

    return df_filled


# 异常值检测与处理函数
def detect_and_handle_outliers(df, numeric_cols, threshold=1.5, strategy='cap', dataset_name=''):
    """
    检测并处理数值型特征的异常值
    参数:
        df: 输入DataFrame
        numeric_cols: 数值型特征列表
        threshold: IQR倍数阈值
        strategy: 处理策略 ('remove', 'cap', 'log', 'zscore')
        dataset_name: 数据集名称（用于输出信息）
    返回:
        处理后的DataFrame, 异常值信息
    """
    df_processed = df.copy()
    outliers_info = {}

    print(f"\n[{dataset_name}] 开始异常值检测与处理...")

    # 保存原始数据分布信息
    original_stats = {}
    for col in numeric_cols:
        original_stats[col] = {
            'min': df[col].min(),
            'max': df[col].max(),
            'mean': df[col].mean(),
            'std': df[col].std()
        }

    # 检测并处理每个数值列
    for col in numeric_cols:
        # 跳过非数值列
        if not np.issubdtype(df[col].dtype, np.number):
            continue

        # 计算四分位数和IQR
        Q1 = df[col].quantile(0.25)
        Q3 = df[col].quantile(0.75)
        IQR = Q3 - Q1

        # 计算异常值边界
        lower_bound = Q1 - threshold * IQR
        upper_bound = Q3 + threshold * IQR

        # 识别异常值
        outlier_mask = (df[col] < lower_bound) | (df[col] > upper_bound)
        outlier_count = outlier_mask.sum()
        outlier_percentage = (outlier_count / len(df)) * 100

        # 记录异常值信息
        outliers_info[col] = {
            'count': outlier_count,
            'percentage': outlier_percentage,
            'lower_bound': lower_bound,
            'upper_bound': upper_bound,
            'indices': df.index[outlier_mask].tolist()
        }

        # 处理异常值
        if outlier_count > 0:
            if strategy == 'remove':
                # 删除异常值
                df_processed = df_processed[~outlier_mask]
                print(f" - 列 '{col}': 删除了 {outlier_count} 个异常值 ({outlier_percentage:.2f}%)")

            elif strategy == 'cap':
                # 缩尾处理（封顶）
                df_processed.loc[df_processed[col] < lower_bound, col] = lower_bound
                df_processed.loc[df_processed[col] > upper_bound, col] = upper_bound
                print(f" - 列 '{col}': 封顶处理了 {outlier_count} 个异常值 ({outlier_percentage:.2f}%)")

            elif strategy == 'log' and (df[col] > 0).all():
                # 对数变换（仅适用于正值）
                df_processed[col] = np.log1p(df_processed[col])
                print(f" - 列 '{col}': 应用了对数变换处理 {outlier_count} 个异常值")

            elif strategy == 'zscore':
                # Z-score标准化处理
                z_scores = np.abs(stats.zscore(df_processed[col]))
                outlier_mask = z_scores > threshold
                median_val = df_processed[col].median()
                df_processed.loc[outlier_mask, col] = median_val
                print(f" - 列 '{col}': 使用中位数替换了 {outlier_mask.sum()} 个Z-score异常值")

    # 重置索引（如果删除了行）
    if strategy == 'remove':
        df_processed = df_processed.reset_index(drop=True)

    # 打印处理前后对比
    print(f"\n[{dataset_name}] 异常值处理前后对比:")
    for col in numeric_cols:
        if col in original_stats and col in df_processed:
            new_min = df_processed[col].min()
            new_max = df_processed[col].max()
            new_mean = df_processed[col].mean()
            new_std = df_processed[col].std()

            print(f" - {col}:")
            print(f"   原始范围: [{original_stats[col]['min']:.2f}, {original_stats[col]['max']:.2f}] "
                  f"均值: {original_stats[col]['mean']:.2f} ± {original_stats[col]['std']:.2f}")
            print(f"   处理后范围: [{new_min:.2f}, {new_max:.2f}] "
                  f"均值: {new_mean:.2f} ± {new_std:.2f}")

    return df_processed, outliers_info


# 数据类型转换
def convert_data_types(df, numeric_cols, categorical_cols):
    """转换数据类型"""
    df_converted = df.copy()

    # 转换为数值型
    for col in numeric_cols:
        df_converted[col] = pd.to_numeric(df_converted[col], errors='coerce')

    # 转换为分类型
    for col in categorical_cols:
        df_converted[col] = df_converted[col].astype('category')

    return df_converted


# 分类变量编码
def encode_categorical_variables(df, categorical_cols, encoding_method='one_hot'):
    """编码分类变量"""
    df_encoded = df.copy()

    if encoding_method == 'label':
        le = LabelEncoder()
        for col in categorical_cols:
            df_encoded[col] = le.fit_transform(df_encoded[col])
    elif encoding_method == 'one_hot':
        encoder = OneHotEncoder(sparse_output=False, drop='first')
        encoded_vars = encoder.fit_transform(df_encoded[categorical_cols])
        encoded_df = pd.DataFrame(encoded_vars, columns=encoder.get_feature_names_out(categorical_cols))
        df_encoded = pd.concat([df_encoded.drop(columns=categorical_cols), encoded_df], axis=1)

    return df_encoded


# --- 1. 中风数据集 (stroke.csv) 处理 - 使用KNN填充数值特征 ---
print("=" * 30 + " 1. 中风数据集 (stroke.csv) 处理 " + "=" * 30)

# 读取数据，并将'N/A'和'Unknown'（针对smoking_status）视作缺失值
stroke_df = pd.read_csv('../data/stroke.csv', na_values=['N/A'])
print("\n原始stroke分布：")
print(stroke_df['stroke'].value_counts())

print("\n(1.1) 原始数据信息及样本预览:")
stroke_df.info()
print(stroke_df.head())

print("\n(1.2) 原始缺失值统计:")
print(stroke_df.isnull().sum())

# 删除'Other'性别类别
other_gender_count = stroke_df[stroke_df['gender'] == 'Other'].shape[0]
if other_gender_count > 0:
    stroke_df = stroke_df[stroke_df['gender'] != 'Other']
    print(f"\n[处理] 删除了 {other_gender_count} 行 'gender' 列为 'Other' 的记录。")
print("\n删除Other后stroke分布：")
print(stroke_df['stroke'].value_counts())

# 识别数值型和分类型特征
numeric_cols_stroke = stroke_df.select_dtypes(include=np.number).columns.tolist()
categorical_cols_stroke = stroke_df.select_dtypes(exclude=np.number).columns.tolist()

# 数据类型转换
stroke_df = convert_data_types(stroke_df, numeric_cols_stroke, categorical_cols_stroke)
print("\n数据类型转换后stroke分布：")
print(stroke_df['stroke'].value_counts())

# 使用KNN填充数值特征（包括BMI）
print("\n[处理] 使用KNN填充数值特征...")
stroke_df = knn_impute(stroke_df, numeric_cols_stroke)
print("KNN填充后stroke分布：")
print(stroke_df['stroke'].value_counts())
print(" - 数值特征缺失值已用KNN填充")

# 使用众数填充分类特征
print("\n[处理] 使用众数填充分类特征...")
for col in categorical_cols_stroke:
    if stroke_df[col].isnull().any():
        mode_val = stroke_df[col].mode()[0]
        stroke_df[col].fillna(mode_val, inplace=True)
        print(f" - 分类型 '{col}' 的缺失值已用众数({mode_val})填充。")
print("众数填充后stroke分布：")
print(stroke_df['stroke'].value_counts())

# 异常值检测与处理（使用封顶策略）
numeric_cols_for_outliers = [col for col in numeric_cols_stroke if col != 'stroke']
stroke_df, stroke_outliers_info = detect_and_handle_outliers(
    stroke_df,
    numeric_cols_for_outliers,
    strategy='cap',
    dataset_name='中风数据集'
)

# 分类变量编码
stroke_df = encode_categorical_variables(stroke_df, categorical_cols_stroke, encoding_method='one_hot')
print("编码后stroke分布：")
print(stroke_df['stroke'].value_counts())

print("\n(1.3) 处理后缺失值统计:")
print(stroke_df.isnull().sum())

# 数值型列保留两位小数
float_cols = stroke_df.select_dtypes(include=['float']).columns
stroke_df[float_cols] = stroke_df[float_cols].round(2)
print("\n[格式化] 数值型列已四舍五入保留两位小数")

# 统一年龄列名
if 'age' in stroke_df.columns:
    stroke_df.rename(columns={'age': 'Age'}, inplace=True)

# 保存处理后的数据
stroke_output_path = os.path.join(output_dir, "stroke_processed.csv")
stroke_df.to_csv(stroke_output_path, index=False)
print(f"\n[保存] 处理后的中风数据集已保存到: {stroke_output_path}")
print("\n中风数据集缺失值处理完成。")

# --- 2. 心脏病数据集 (heart.csv) 处理 - 保留原方式（无缺失值）---
print("\n" + "=" * 30 + " 2. 心脏病数据集 (heart.csv) 处理 " + "=" * 30)

# 读取数据，指定'NA'为缺失值标识
heart_df = pd.read_csv('../data/heart.csv', na_values=['NA'])

print("\n(2.1) 原始数据信息及样本预览:")
heart_df.info()
print(heart_df.head())

print("\n(2.2) 原始缺失值统计:")
missing_heart = heart_df.isnull().sum()
print(missing_heart)

# 检查是否存在缺失值
if missing_heart.sum() == 0:
    print("\n[结论] 心脏病数据集没有发现'NA'标识的缺失值。")
else:
    print("\n[处理] 心脏病数据集中发现缺失值，将进行处理...")
    for col in heart_df.columns:
        if heart_df[col].isnull().any():
            if pd.api.types.is_numeric_dtype(heart_df[col]):
                median_val = heart_df[col].median()
                heart_df[col].fillna(median_val, inplace=True)
                print(f" - 数值列 '{col}' 的缺失值已用中位数({median_val})填充.")
            else:
                mode_val = heart_df[col].mode()[0]
                heart_df[col].fillna(mode_val, inplace=True)
                print(f" - 分类型 '{col}' 的缺失值已用众数({mode_val})填充.")
    print("\n(2.3) 处理后缺失值统计:")
    print(heart_df.isnull().sum())

# 异常值检测与处理（使用封顶策略）
numeric_cols_heart = heart_df.select_dtypes(include=np.number).columns.tolist()
categorical_cols_heart = heart_df.select_dtypes(exclude=np.number).columns.tolist()

# 排除目标变量
if 'target' in numeric_cols_heart:
    numeric_cols_heart.remove('target')

heart_df, heart_outliers_info = detect_and_handle_outliers(
    heart_df,
    numeric_cols_heart,
    strategy='cap',
    dataset_name='心脏病数据集'
)

# 数值型列保留两位小数
float_cols = heart_df.select_dtypes(include=['float']).columns
heart_df[float_cols] = heart_df[float_cols].round(2)
print("\n[格式化] 数值型列已四舍五入保留两位小数")

# 保存处理后的数据
heart_output_path = os.path.join(output_dir, "heart_processed.csv")
heart_df.to_csv(heart_output_path, index=False)
print(f"\n[保存] 处理后的心脏病数据集已保存到: {heart_output_path}")
print("\n心脏病数据集缺失值处理完成。")

# --- 3. 肝硬化数据集 (cirrhosis.csv) 处理 - 使用KNN填充数值特征 ---
print("\n" + "=" * 30 + " 3. 肝硬化数据集 (cirrhosis.csv) 处理 " + "=" * 30)

# 读取数据，指定 'NA' 为缺失值标识
cirrhosis_df = pd.read_csv('../data/cirrhosis.csv', na_values=['NA'])

print("\n(3.1) 原始数据信息及样本预览:")
cirrhosis_df.info()
print(cirrhosis_df.head())

print("\n(3.2) 原始缺失值统计:")
print(cirrhosis_df.isnull().sum())

# 删除目标变量'Stage'缺失的行
initial_rows = len(cirrhosis_df)
cirrhosis_df.dropna(subset=['Stage'], inplace=True)
print(f"\n[处理] 'Stage' 列存在缺失值的行已被删除，共删除 {initial_rows - len(cirrhosis_df)} 行。")

# 识别数值型和分类型特征
numeric_cols_cirrhosis = cirrhosis_df.select_dtypes(include=np.number).columns.tolist()
categorical_cols_cirrhosis = cirrhosis_df.select_dtypes(exclude=np.number).columns.tolist()

# 数据类型转换
cirrhosis_df = convert_data_types(cirrhosis_df, numeric_cols_cirrhosis, categorical_cols_cirrhosis)

# 使用KNN填充数值特征
print("\n[处理] 使用KNN填充数值特征...")
cirrhosis_df = knn_impute(cirrhosis_df, numeric_cols_cirrhosis)
print(" - 数值特征缺失值已用KNN填充")

# 使用众数填充分类特征
print("\n[处理] 使用众数填充分类特征...")
for col in categorical_cols_cirrhosis:
    if cirrhosis_df[col].isnull().any() and col != 'Stage':
        mode_val = cirrhosis_df[col].mode()[0]
        cirrhosis_df[col].fillna(mode_val, inplace=True)
        print(f" - 分类型 '{col}' 的缺失值已用众数({mode_val})填充。")

# 异常值检测与处理（使用封顶策略）
if 'Stage' in numeric_cols_cirrhosis:
    numeric_cols_cirrhosis.remove('Stage')

cirrhosis_df, cirrhosis_outliers_info = detect_and_handle_outliers(
    cirrhosis_df,
    numeric_cols_cirrhosis,
    strategy='cap',
    dataset_name='肝硬化数据集'
)

# 分类变量编码
cirrhosis_df = encode_categorical_variables(cirrhosis_df, categorical_cols_cirrhosis, encoding_method='one_hot')

print("\n(3.3) 处理后缺失值统计:")
print(cirrhosis_df.isnull().sum())

# 数值型列保留两位小数
float_cols = cirrhosis_df.select_dtypes(include=['float']).columns
cirrhosis_df[float_cols] = cirrhosis_df[float_cols].round(2)
print("\n[格式化] 数值型列已四舍五入保留两位小数")

# 保存处理后的数据
cirrhosis_output_path = os.path.join(output_dir, "cirrhosis_processed.csv")
cirrhosis_df.to_csv(cirrhosis_output_path, index=False)
print(f"\n[保存] 处理后的肝硬化数据集已保存到: {cirrhosis_output_path}")
print("\n肝硬化数据集缺失值处理完成。")